#ifndef LINEARSYSTEM
#define LINEARSYSTEM

#include <stdlib.h>
#include <math.h>

#define SOLVE_GAUSSIAN 0 	
#define SOLVE_JACOBI 1 
#define SOLVE_STEEPESTDESCENT 2 
#define SOLVE_CONJUGATEGRADIENT 3 

/**
 * Class for solving linear system of equations
 */
class LinearSystemSolver
{
public:
	LinearSystemSolver ( ) { } ;
	
	/**
	 * Entry method to solve linear system of equations (force to use double solver)
	 */
	template <class T>
	static void solveD ( int dim, T** A, T* b, T* res, int type )
	{
		int i, j ;
		
		double** dA = new double *[dim] ;
		double* db = new double[dim] ;
		double* dres = new double[dim] ;
		for ( i = 0 ; i < dim ; i ++ )
		{
			dA[i] = new double [ dim ] ;
		}
		
		for ( i = 0 ; i < dim ; i ++ )
			for ( j = 0 ; j < dim ; j ++ )
			{
				dA[i][j] = (double) A[i][j] ;	
			}
			for ( i = 0 ; i < dim ; i ++ )
			{
				db[i] = (double) b[i] ;
			}
			
			for ( i = 0 ; i < dim ; i ++ )
			{
				dres[i] = (double) res[i] ;
			}
			
			solve ( dim, dA, db, dres, type ) ;
			
			for ( i = 0 ; i < dim ; i ++ )
			{
				res[i] = (float) dres[i] ;
			}
			
			for ( i = 0 ; i < dim ; i ++ )
			{
				delete[] dA[i] ;
			}
			delete[] dA ;
			delete[] db ;
			delete[] dres ;
			
	}	

	/**
	 * Method to retrieve residue
	 */
	template <class T>
	static T getResidue ( int dim1, int dim2, T** A, T* b, T* res ) 
	{
		
		T* r = new T[dim2] ;
		matrixMultiply( dim1, dim2, A, dim2, res, r, 0 ) ;
		vectorAddition( dim2, b, r, r, 0 ) ;
		T rvalue = vectorMultiply( dim2, r, r ) ;
		delete[] r ;
		
		return rvalue ;
	}
	
	/**
	 * Entry method to solve linear system of equations
	 */
	template <class T>
	static void solve ( int dim, T** A, T* b, T* res, int type ) 
	{
		/// Next, solve the system
		switch ( type )
		{
		case SOLVE_GAUSSIAN :
			solveGaussian ( dim, A, b, res ) ;
			break ;
		case SOLVE_JACOBI :
			solveJacobi ( dim, A, b, res ) ;
			break ;
		case SOLVE_STEEPESTDESCENT :
			solveSD ( dim, A, b, res ) ;
			break ;
		case SOLVE_CONJUGATEGRADIENT :
			solveCG ( dim, A, b, res ) ;
			break ;
		}
	}	
	
	/**
	 * Gaussian elimination
	 */
	template <class T>
	static void solveGaussian ( int dim, T** A, T* b, T* rvalue ) 
	{
		int i, j, k ;
		int len = dim ;
		int* order = new int[len] ;
		int* mask = new int[len] ;
		for ( i = 0 ; i < len ; i ++ )
		{
			mask[i] = 0 ;
			order[i] = -1 ;
		}
		
		/// Elimination
		T epsilon = (T) 0.000000001 ;
		for ( i = 0 ; i < len ; i ++ )
		{
			/// Search for the largest coefficient in the ith column
			T biggest = 0, temp ;
			int bigindex = -1 ;
			for ( j = 0 ; j < len ; j ++ )
			{
				if ( !mask[j] )	
				{
					if ( (temp = (T) fabs(A[j][i])) > biggest )
					{
						biggest = temp ;
						bigindex = j;
					}
				}
			}
			
			if ( biggest < epsilon )
			{
				printf("Gaussian elimination: coefficient matrix is singular.\n");
				// return ;
			}
			
			mask[bigindex] = 1 ;
			order[i] = bigindex ;
			
			/// Eliminate other rows using row j
			for ( j = 0 ; j < len ; j ++ )
			{
				if 	( ! mask[j] )
				{
					T rate = A[j][i] / A[bigindex][i] ;
					for ( k = i ; k < len ; k ++ )
					{
						A[j][k] -= A[bigindex][k] * rate ;
					}
					
					b[j] -= b[bigindex] * rate ;
				}
			}
		}
		
		/// Solution
		for ( i = len - 1 ; i >= 0 ; i -- )
		{
			T sum = 0 ;
			for ( k = i + 1 ; k < len ; k ++ )
			{
				sum += A[order[i]][k] * rvalue[k] ;	
			}
			rvalue[i] = ( b[order[i]] - sum ) / A[order[i]][i] ;
		}

		T error = getResidue ( dim, dim, A, b, rvalue ) ;
		
		printf("Error: %g\n", error ) ;
		
		delete[] order ;
		delete[] mask ;
	}
	
	/**
	 * Jacobi Iteration
	 */
	template <class T>
	static void solveJacobi ( int dim, T** A, T* b, T* rvalue )
	{
		int i, j, len = dim ;
		T** D = new T * [len] ;
		T** E = new T * [len] ;
		T** B = new T * [len] ;

		for ( i = 0 ; i < len ; i ++ )
		{
			D[i] = new T[len] ;
			E[i] = new T[len] ;
			B[i] = new T[len] ;
		}
		
		for ( i = 0 ; i < len ; i ++ )
			for ( j = 0 ; j < len ; j ++ )
			{
				if ( i == j )
				{
					D[i][j] = 1 / A[i][j] ;
				}
				else 
				{
					E[i][j] = A[i][j] ;
				}
			}
			
			T* z = new T[len];
			
			matrixMultiply( dim, dim, D, dim, dim, E, B, 0, 0 ) ;
			matrixMultiply( dim, dim, D, dim, b, z, 0 ) ;
			
			T epsilon = (T) 0.000001 ;
			T error = getResidue( dim, dim, A, b, rvalue ) ;
			T* tx = new T[len];
			while ( error > epsilon )
			{
				printf("Error: %g\n", error ) ;

				/// Iterate	
				matrixMultiply( dim, dim, B, dim, rvalue, tx, 0 ) ;
				vectorAddition( dim, z, tx, rvalue, 0 ) ;
				
				error = getResidue( dim, dim, A, b, rvalue ) ;
			}

		for ( i = 0 ; i < len ; i ++ )
		{
			delete[] D[i] ;
			delete[] E[i] ;
			delete[] B[i] ;
		}

		delete[] tx ;
		delete[] z ;
		delete[] D ;
		delete[] E ;
		delete[] B ;
	}
																  
	/**
	 * Steepest Descent
	 */
	template <class T>
	static void solveSD ( int dim, T** A, T* b, T* rvalue )
	{
		int len = dim ;
		T* r = new T[ len ] ;
		T* r1 = new T[ len ] ;
		T* r2 = new T[ len ] ;
		T alpha ;
		
		/// Setup r_0
		matrixMultiply( dim, dim, A, dim, rvalue, r, 0 ) ;
		vectorAddition( dim, b, r, r, 0 ) ;
		
		/// Iterate
		T epsilon = (T) 0.0001 ;
		T error = getResidue( dim, dim, A, b, rvalue ) ;
		// while ( error > epsilon )
		for ( int i = 0 ; i < 10000 ; i ++ )
		{
			printf("Error: %g\n", error ) ;
			
			// a_i
			matrixMultiply( dim, dim, A, dim, r, r1, 0 ) ;
			alpha = vectorMultiply ( dim, r, r ) / vectorMultiply ( dim, r, r1 ) ;
			
			// x_i+1
			vectorScale( dim, r, alpha, r2, 1 ) ;
			vectorAddition ( dim, rvalue, r2, rvalue, 1 ) ;
			
			// r_i+1
			vectorScale( dim, r1, alpha, r2, 1 ) ;
			vectorAddition ( dim, r, r2, r, 0 ) ;
			
			error = getResidue ( dim, dim, A, b, rvalue ) ;
		}

		delete[] r ;
		delete[] r1 ;
		delete[] r2 ;
	}
																  
	/**
	 * Conjugate Gradient
	 */
	template <class T>
	static void solveCG ( int dim, T** A, T* b, T* rvalue ) 
	{
		int len = dim ;
		T* d = new T[ len ] ;
		T* r = new T[ len ] ;
		T* r1 = new T[ len ] ;
		T* d1 = new T[ len ] ;
		T* d2 = new T[ len ] ;
		T alpha, beta ;
		
		/// Setup r_0, d_0
		matrixMultiply( dim, dim, A, dim, rvalue, r, 0 ) ;
		vectorAddition( dim, b, r, r, 0 ) ;
		vectorScale( dim, r, (T) 1, d, 1 ) ;
		
		/// Iterate
		T epsilon = (T) 0.000000001 ;
		T error = getResidue( dim, dim, A, b, rvalue ) ;
		T initError = error ;
		
		int i = 0 ;
		
		//while ( i < 5 )
	   while ( error > epsilon * epsilon * initError && i < 1000 ) 
		{
			error = getResidue ( dim, dim, A, b, rvalue ) ;
			// System.out.println("Error: " + error ) ;
			// System.out.println("Residue: " + matrixMultiply( r, r ) ) ;
			
			// alpha_i
			matrixMultiply( dim, dim, A, dim, d, d1, 0 ) ;
			// System.out.println("d: " + matrixMultiply ( d, d ) ) ;
			// System.out.println("d1: " + matrixMultiply ( d1, d1 ) ) ;
			
			alpha = vectorMultiply ( dim, r, r ) / vectorMultiply ( dim, d, d1 ) ;
			
			// x_i+1
			vectorScale( dim, d, alpha, d2, 1 ) ;
			vectorAddition ( dim, rvalue, d2, rvalue, 1 ) ;
			
			// r_i+1
			vectorScale ( dim, r, (T) 1, r1, 1 ) ;
			//if ( i % 50 == 0 )
			//{
			//	vectorScale( alpha, d1, d2, 1 ) ;
			//	vectorAddition ( r, d2, r, 0 ) ;
			//}
			//else
			{
				matrixMultiply( dim, dim, A, dim, rvalue, d2, 0 ) ;
				vectorAddition ( dim, b, d2, r, 0 ) ;
			}
			
			// beta_i+1
			beta = vectorMultiply ( dim, r, r ) / vectorMultiply ( dim, r1, r1 ) ;
			
			// d_i+1
			vectorScale( dim, d, beta, d2, 1 ) ;
			vectorAddition ( dim, r, d2, d, 1 ) ;
			
			i ++ ;
		}
		
		printf("Residue: %g\n", vectorMultiply( dim, r, r ) ) ;

		delete[] d;
		delete[] d1 ;
		delete[] d2 ;
		delete[] r ;
		delete[] r1 ;
	}
	
	/**
	 * Matrix operations
	 * Note 1: the correctness of matrix dimensions need to be verified.
	 * Note 2: Every vector is a row vector. Every matrix has shape [numRows][numColumns] .
	 */
	template <class T>
	static void matrixScale ( int dim1, int dim2, T** A, T s, T** B )
	{
		for ( int i = 0 ; i < dim1 ; i ++ )
		{
			for ( int j = 0 ; j < dim2 ; j ++ )
			{
				B[i][j] = A[i][j] * s ;
			}
		}
	}	

	template <class T>
	static void vectorScale ( int dim, T* a, T s, T* b, int multiply )
	{
		if ( multiply )
		{
			for ( int i = 0 ; i < dim ; i ++ )
			{
				b[i] = a[i] * s ;
			}
		}
		else
		{
			for ( int i = 0 ; i < dim ; i ++ )
			{
				b[i] = a[i] / s ;
			}
		}
	}
	
	template <class T>
	static void matrixAddition ( int dim1, int dim2, T** A, T** B, T** C, int addition ) 
	{
		for ( int i = 0 ; i < dim1 ; i ++ )
		{
			for ( int j = 0 ; j < dim2 ; j ++ )
			{
				if ( addition )
				{
					C[i][j] = A[i][j] + B[i][j] ;	
				}
				else 
				{
					C[i][j] = A[i][j] - B[i][j] ;	
				}
			}
		}
	}
	
	template <class T>
	static void vectorAddition ( int dim, T* a, T* b, T* res, int addition )
	{
		for ( int i = 0 ; i < dim ; i ++ )
		{
			if ( addition )
			{
				res[i] = a[i] + b[i] ;	
			}
			else 
			{
				res[i] = a[i] - b[i] ;	
			}
		}
	}
	
	template <class T>
	static float vectorMultiply ( int dim, T* a, T* b )
	{
		T rvalue = 0 ;
		for ( int i = 0 ; i < dim ; i ++ )
		{
			rvalue += b[i] * a[i] ;
		}
		
		return rvalue ;
	}
	// dim1: length of a
	// dim2: length of b
	template <class T>
	static void matrixMultiply ( int dim1, T* a, int dim2, T* b, T** res )
	{
		T** ta = new T * [1] ;
		ta[0] = a ;
		T** tb = new T * [1] ;
		tb[0] = b ;
		
		matrixMultiply( 1, dim1, ta, 1, dim2, tb, res, 1, 0 ) ;

		delete[] ta;
		delete[] tb ;
	}

	// dim1: length of a
	// dim2, dim3: length/width of B
	template <class T>
	static void matrixMultiply ( int dim1, T* a, int dim2, int dim3, T** B, T* res, int transposeB )
	{
		T** ta = new T * [1] ;
		ta[0] = a ;
		T** tres = new T * [1] ;
		tres[0] = res ;
		
		matrixMultiply( 1, dim1, ta, dim2, dim3, B, tres, 0, transposeB ) ;

		delete[] ta;
		delete[] tres ;
	}
		
	// dim1, dim2: length/width of A
	// dim3: length of b
	template <class T>
	static void matrixMultiply ( int dim1, int dim2, T** A, int dim3, T* b, T* res, int transposeA )
	{
		T** tb = new T * [1] ;
		tb[0] = b ;
		T** tres = new T * [1] ;
		tres[0] = res ;
		
		matrixMultiply( 1, dim3, tb, dim1, dim2, A, tres, 0, !transposeA ) ;
		delete[] tb ;
		delete[] tres ;
	}
	
	template <class T>
	static void matrixMultiply ( int dim1, int dim2, T** A, int dim3, int dim4, T** B, T** res, int transposeA, int transposeB )
	{
		int i, j, k ;
		if ( !transposeA && !transposeB )
		{
			for ( i = 0 ; i < dim1 ; i ++ )
				for ( j = 0 ; j < dim4 ; j ++ )
				{
					res[i][j] = 0 ;
					for ( k = 0 ; k < dim3 ; k ++ )
					{
						res[i][j] += A[i][k] * B[k][j] ;	
					}
				}
		}
		else if ( transposeA && !transposeB )
		{
			for ( i = 0 ; i < dim2 ; i ++ )
				for ( j = 0 ; j < dim4 ; j ++ )
				{
					res[i][j] = 0 ;
					for ( k = 0 ; k < dim3 ; k ++ )
					{
						res[i][j] += A[k][i] * B[k][j] ;	
					}
				}
		}
		else if ( !transposeA && transposeB )
		{
			for ( i = 0 ; i < dim1 ; i ++ )
				for ( j = 0 ; j < dim3 ; j ++ )
				{
					res[i][j] = 0 ;
					for ( k = 0 ; k < dim4 ; k ++ )
					{
						res[i][j] += A[i][k] * B[j][k] ;	
					}
				}
		}
		else if ( transposeA && transposeB )
		{
			for ( i = 0 ; i < dim2 ; i ++ )
				for ( j = 0 ; j < dim3 ; j ++ )
				{
					res[i][j] = 0 ;
					for ( k = 0 ; k < dim1 ; k ++ )
					{
						res[i][j] += A[k][i] * B[j][k] ;	
					}
				}
		}
	}
	
	template <class T>
	static T matrixTrace ( int dim, T** A )
	{
		T tr = 0 ;	
		
		for ( int i = 0 ; i < dim ; i ++ )
		{
			tr += A[i][i] ;	
		}
		
		return tr ;
	}

	template <class T>
	static void printVector( int dim, T* a ) 
	{
		for ( int i = 0 ; i < dim ; i ++ )
		{
			printf("%g ", a[i]);
		}
		printf("\n") ;
	}
	
	template <class T>
	static void printMatrix( int dim1, int dim2, T** A ) 
	{
		for ( int i = 0 ; i < dim1 ; i ++ )
		{
			printVector( dim2, A[i] ) ;
		}
		printf("\n");
	}

};


#endif